# algorithms.py
import heapq
from collections import deque, defaultdict
import math

# 1. Dijkstra 最短路径
def dijkstra(graph, start, end):
    dist = [float('inf')] * graph.V
    dist[start] = 0
    parent = [-1] * graph.V
    visited = set()
    heap = [(0, start)]

    while heap:
        d, u = heapq.heappop(heap)
        if u in visited:
            continue
        visited.add(u)
        yield 'explore', u, None
        if u == end:
            break
        for v in range(graph.V):
            if graph.graph[u][v] < float('inf'):
                alt = dist[u] + graph.graph[u][v]
                if alt < dist[v]:
                    dist[v] = alt
                    parent[v] = u
                    heapq.heappush(heap, (alt, v))
                    yield 'update', v, u
    # 回溯路径
    path = []
    if dist[end] != float('inf'):
        cur = end
        while cur != -1:
            path.append(cur)
            cur = parent[cur]
        path.reverse()
    yield 'path', path, None

# 2. Bellman-Ford（支持负权）
def bellman_ford(graph, start, end):
    dist = [float('inf')] * graph.V
    dist[start] = 0
    parent = [-1] * graph.V

    for i in range(graph.V - 1):
        updated = False
        for u in range(graph.V):
            for v in range(graph.V):
                if graph.graph[u][v] < float('inf'):
                    if dist[u] + graph.graph[u][v] < dist[v]:
                        dist[v] = dist[u] + graph.graph[u][v]
                        parent[v] = u
                        updated = True
                        yield 'update', v, u
        if not updated:
            break

    # 检查负环
    for u in range(graph.V):
        for v in range(graph.V):
            if graph.graph[u][v] < float('inf') and dist[u] + graph.graph[u][v] < dist[v]:
                yield 'negative_cycle', None, None
                return

    path = []
    if dist[end] != float('inf'):
        cur = end
        while cur != -1:
            path.append(cur)
            cur = parent[cur]
        path.reverse()
    yield 'path', path, None

# 3. A* 算法（需启发函数）
def heuristic(a, b, pos):
    return math.sqrt((pos[a][0] - pos[b][0])**2 + (pos[a][1] - pos[b][1])**2)

def a_star(graph, start, end, pos):
    open_set = [(0, start)]
    g_score = [float('inf')] * graph.V
    g_score[start] = 0
    f_score = [float('inf')] * graph.V
    f_score[start] = heuristic(start, end, pos)
    parent = [-1] * graph.V
    visited = set()

    while open_set:
        current_f, u = heapq.heappop(open_set)
        if u == end:
            break
        if u in visited:
            continue
        visited.add(u)
        yield 'explore', u, None

        for v in range(graph.V):
            if graph.graph[u][v] < float('inf'):
                tentative_g = g_score[u] + graph.graph[u][v]
                if tentative_g < g_score[v]:
                    parent[v] = u
                    g_score[v] = tentative_g
                    f_score[v] = tentative_g + heuristic(v, end, pos)
                    heapq.heappush(open_set, (f_score[v], v))
                    yield 'update', v, u

    path = []
    cur = end
    while cur != -1 and parent[cur] != -1 or cur == start:
        path.append(cur)
        if cur == start:
            break
        cur = parent[cur]
    else:
        path = []
    path.reverse()
    yield 'path', path, None

# 4. Floyd-Warshall（所有点对最短路径）
def floyd_warshall(graph, start, end):
    dist = graph.graph.copy()
    parent = [[-1 for _ in range(graph.V)] for _ in range(graph.V)]

    for k in range(graph.V):
        for i in range(graph.V):
            for j in range(graph.V):
                if dist[i][k] + dist[k][j] < dist[i][j]:
                    dist[i][j] = dist[i][k] + dist[k][j]
                    parent[i][j] = k
        yield 'step', k, None  # 显示每一轮

    # 提取路径
    def get_path(i, j):
        if dist[i][j] == float('inf'):
            return []
        if parent[i][j] == -1:
            return [i, j] if i != j else [i]
        k = parent[i][j]
        return get_path(i, k)[:-1] + get_path(k, j)

    path = get_path(start, end)
    yield 'path', path, None

# 5. DFS 枚举所有路径
def dfs_all_paths(graph, start, end):
    path = []
    visited = [False] * graph.V

    def dfs(u):
        visited[u] = True
        path.append(u)
        if u == end:
            yield 'found', path.copy(), None
        else:
            for v in range(graph.V):
                if graph.graph[u][v] < float('inf') and not visited[v]:
                    yield from dfs(v)
        path.pop()
        visited[u] = False

    yield from dfs(start)

# 6. SPFA（队列优化 Bellman-Ford）
def spfa(graph, start, end):
    dist = [float('inf')] * graph.V
    dist[start] = 0
    parent = [-1] * graph.V
    in_queue = [False] * graph.V
    queue = deque([start])
    in_queue[start] = True

    while queue:
        u = queue.popleft()
        in_queue[u] = False
        yield 'explore', u, None
        for v in range(graph.V):
            if graph.graph[u][v] < float('inf'):
                new_dist = dist[u] + graph.graph[u][v]
                if new_dist < dist[v]:
                    dist[v] = new_dist
                    parent[v] = u
                    if not in_queue[v]:
                        queue.append(v)
                        in_queue[v] = True
                    yield 'update', v, u

    path = []
    if dist[end] != float('inf'):
        cur = end
        while cur != -1:
            path.append(cur)
            cur = parent[cur]
        path.reverse()
    yield 'path', path, None
